use crate::ast::*;
use crate::error::{CompilerError, Result};
use crate::llvm_backend::LLVMBackend;
use inkwell::values::BasicValueEnum;

impl<'ctx> LLVMBackend<'ctx> {


    /// Generate map() method - transforms Some(x) to Some(f(x)), None stays None
    pub(super) fn generate_option_map(
        &mut self,
        option_value: BasicValueEnum<'ctx>,
        option_type: &crate::ast::types::OptionType,
        arguments: &[Expression],
        location: &crate::error::Location,
    ) -> Result<BasicValueEnum<'ctx>> {
        if arguments.len() != 1 {
            return Err(CompilerError::type_error(
                location.line,
                location.column,
                &format!(
                    "Option.map() expects 1 argument (function), got {}",
                    arguments.len()
                ),
            ));
        }

        let func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("Cannot create blocks outside function"))?;

        let tag = self.extract_option_tag(option_value)?;
        let zero = self.context.i32_type().const_zero();
        let is_some =
            self.builder
                .build_int_compare(inkwell::IntPredicate::EQ, tag, zero, "map_check")?;

        // Create blocks: some_block (apply function), none_block (keep None)
        let some_block = self.context.append_basic_block(func, "map_some");
        let none_block = self.context.append_basic_block(func, "map_none");
        let continue_block = self.context.append_basic_block(func, "map_continue");

        self.builder
            .build_conditional_branch(is_some, some_block, none_block)?;

        // Some branch: extract Some value, call function, create new Option::Some
        self.builder.position_at_end(some_block);
        let some_value = self.extract_option_some_value(option_value, option_type)?;

        // Store some_value in a temporary variable so we can use it in function call
        let some_type = self.nrc_type_to_llvm_type(&Some(*option_type.some_type.clone()))?;
        let some_temp_alloca = self.builder.build_alloca(some_type, "some_temp")?;
        self.builder.build_store(some_temp_alloca, some_value)?;

        // Create a temporary variable name for the some value
        let some_temp_var = format!("__map_some_{}", location.line);

        // Store the alloca in variable_map temporarily
        let old_some_value = self
            .variable_map
            .insert(some_temp_var.clone(), some_temp_alloca.into());
        let old_some_type = self
            .variable_type_map
            .insert(some_temp_var.clone(), *option_type.some_type.clone());

        // Create a call expression: f(some_temp_var)
        let func_expr = arguments[0].clone();
        let some_var_expr = Expression::Variable(some_temp_var.clone());
        let call_expr = Expression::Call(CallExpr {
            callee: Box::new(func_expr),
            arguments: vec![some_var_expr],
            type_args: Vec::new(),
            resolved: None,
            location: *location,
        });

        // Generate the function call
        let transformed_value = self.generate_expression(&call_expr)?;

        // Infer the transformed value's type from the function return type
        let transformed_type = match self.infer_expression_type(&call_expr) {
            Some(t) => t,
            _ => *option_type.some_type.clone(), // Fallback to Some type
        };

        // Restore variable map
        if let Some(old) = old_some_value {
            self.variable_map.insert(some_temp_var.clone(), old);
        } else {
            self.variable_map.remove(&some_temp_var);
        }
        if let Some(old) = old_some_type {
            self.variable_type_map.insert(some_temp_var, old);
        } else {
            self.variable_type_map.remove(&some_temp_var);
        }

        // Create new Option::Some with transformed value
        let new_option_type = crate::ast::types::OptionType {
            some_type: Box::new(transformed_type.clone()),
            location: *location,
        };
        let new_some_result =
            self.create_option_value(0, Some(transformed_value), &new_option_type)?;

        self.builder.build_unconditional_branch(continue_block)?;

        // None branch: create new Option::None
        self.builder.position_at_end(none_block);
        let new_none_result = self.create_option_value(1, None, option_type)?;
        self.builder.build_unconditional_branch(continue_block)?;

        // Continue block: return either new Some Option or new None Option
        self.builder.position_at_end(continue_block);
        let option_llvm_type = self.get_option_llvm_type(&new_option_type)?;
        let phi = self.builder.build_phi(option_llvm_type, "map_result")?;
        phi.add_incoming(&[
            (&new_some_result, some_block),
            (&new_none_result, none_block),
        ]);
        Ok(phi.as_basic_value())
    }


    /// Generate and_then() method - chains Option operations
    pub(super) fn generate_option_and_then(
        &mut self,
        option_value: BasicValueEnum<'ctx>,
        option_type: &crate::ast::types::OptionType,
        arguments: &[Expression],
        location: &crate::error::Location,
    ) -> Result<BasicValueEnum<'ctx>> {
        if arguments.len() != 1 {
            return Err(CompilerError::type_error(
                location.line,
                location.column,
                &format!(
                    "Option.and_then() expects 1 argument (function), got {}",
                    arguments.len()
                ),
            ));
        }

        let func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("Cannot create blocks outside function"))?;

        let tag = self.extract_option_tag(option_value)?;
        let zero = self.context.i32_type().const_int(0, false);
        let is_some = self.builder.build_int_compare(
            inkwell::IntPredicate::EQ,
            tag,
            zero,
            "and_then_check",
        )?;

        // Create blocks: some_block (call function), none_block (keep None)
        let some_block = self.context.append_basic_block(func, "and_then_some");
        let none_block = self.context.append_basic_block(func, "and_then_none");
        let continue_block = self.context.append_basic_block(func, "and_then_continue");

        self.builder
            .build_conditional_branch(is_some, some_block, none_block)?;

        // Some branch: extract Some value, call function (which returns Option), return that Option
        self.builder.position_at_end(some_block);
        let some_value = self.extract_option_some_value(option_value, option_type)?;

        // Store some_value in a temporary variable
        let some_type = self.nrc_type_to_llvm_type(&Some(*option_type.some_type.clone()))?;
        let some_temp_alloca = self.builder.build_alloca(some_type, "some_temp")?;
        self.builder.build_store(some_temp_alloca, some_value)?;

        let some_temp_var = format!("__and_then_some_{}", location.line);
        let old_some_value = self
            .variable_map
            .insert(some_temp_var.clone(), some_temp_alloca.into());
        let old_some_type = self
            .variable_type_map
            .insert(some_temp_var.clone(), *option_type.some_type.clone());

        // Create call expression: f(some_temp_var)
        let func_expr = arguments[0].clone();
        let some_var_expr = Expression::Variable(some_temp_var.clone());
        let call_expr = Expression::Call(CallExpr {
            callee: Box::new(func_expr),
            arguments: vec![some_var_expr],
            type_args: Vec::new(),
            resolved: None,
            location: *location,
        });

        // Generate the function call (returns Option<U>)
        let chained_option = self.generate_expression(&call_expr)?;

        // Infer the return type
        let chained_option_type = match self.infer_expression_type(&call_expr) {
            Some(Type::Option(opt_type)) => opt_type,
            _ => {
                return Err(CompilerError::internal(
                    "and_then() function must return Option type",
                ))
            }
        };

        // Restore variable map
        if let Some(old) = old_some_value {
            self.variable_map.insert(some_temp_var.clone(), old);
        } else {
            self.variable_map.remove(&some_temp_var);
        }
        if let Some(old) = old_some_type {
            self.variable_type_map.insert(some_temp_var, old);
        } else {
            self.variable_type_map.remove(&some_temp_var);
        }

        self.builder.build_unconditional_branch(continue_block)?;

        // None branch: create new Option::None with same type as chained option
        self.builder.position_at_end(none_block);
        let new_none_result = self.create_option_value(1, None, &chained_option_type)?;
        self.builder.build_unconditional_branch(continue_block)?;

        // Continue block: return either chained Option or None
        self.builder.position_at_end(continue_block);
        let option_llvm_type = self.get_option_llvm_type(&chained_option_type)?;
        let phi = self
            .builder
            .build_phi(option_llvm_type, "and_then_result")?;
        phi.add_incoming(&[
            (&chained_option, some_block),
            (&new_none_result, none_block),
        ]);
        Ok(phi.as_basic_value())
    }


    /// Generate or_else() method - provides fallback Option if None
    pub(super) fn generate_option_or_else(
        &mut self,
        option_value: BasicValueEnum<'ctx>,
        option_type: &crate::ast::types::OptionType,
        arguments: &[Expression],
        location: &crate::error::Location,
    ) -> Result<BasicValueEnum<'ctx>> {
        if arguments.len() != 1 {
            return Err(CompilerError::type_error(
                location.line,
                location.column,
                &format!(
                    "Option.or_else() expects 1 argument (function), got {}",
                    arguments.len()
                ),
            ));
        }

        let func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("Cannot create blocks outside function"))?;

        let tag = self.extract_option_tag(option_value)?;
        let zero = self.context.i32_type().const_int(0, false);
        let is_some = self.builder.build_int_compare(
            inkwell::IntPredicate::EQ,
            tag,
            zero,
            "or_else_check",
        )?;

        // Create blocks: some_block (keep Some), none_block (call function)
        let some_block = self.context.append_basic_block(func, "or_else_some");
        let none_block = self.context.append_basic_block(func, "or_else_none");
        let continue_block = self.context.append_basic_block(func, "or_else_continue");

        self.builder
            .build_conditional_branch(is_some, some_block, none_block)?;

        // Some block: keep the original Option
        self.builder.position_at_end(some_block);
        self.builder.build_unconditional_branch(continue_block)?;

        // None block: call function to get fallback Option
        self.builder.position_at_end(none_block);
        let call_expr = Expression::Call(CallExpr {
            callee: Box::new(arguments[0].clone()),
            arguments: vec![],
            type_args: Vec::new(),
            resolved: None,
            location: *location,
        });
        let fallback_option = self.generate_expression(&call_expr)?;
        self.builder.build_unconditional_branch(continue_block)?;

        // Continue block: return either original Some or fallback Option
        self.builder.position_at_end(continue_block);
        let option_llvm_type = self.get_option_llvm_type(option_type)?;
        let phi = self.builder.build_phi(option_llvm_type, "or_else_result")?;
        phi.add_incoming(&[(&option_value, some_block), (&fallback_option, none_block)]);
        Ok(phi.as_basic_value())
    }


    /// Generate filter() method - filters Some value by predicate
    pub(super) fn generate_option_filter(
        &mut self,
        option_value: BasicValueEnum<'ctx>,
        option_type: &crate::ast::types::OptionType,
        arguments: &[Expression],
        location: &crate::error::Location,
    ) -> Result<BasicValueEnum<'ctx>> {
        if arguments.len() != 1 {
            return Err(CompilerError::type_error(
                location.line,
                location.column,
                &format!(
                    "Option.filter() expects 1 argument (predicate function), got {}",
                    arguments.len()
                ),
            ));
        }

        let func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("Cannot create blocks outside function"))?;

        let tag = self.extract_option_tag(option_value)?;
        let zero = self.context.i32_type().const_int(0, false);
        let is_some =
            self.builder
                .build_int_compare(inkwell::IntPredicate::EQ, tag, zero, "filter_check")?;

        // Create blocks: some_block (check predicate), none_block (keep None)
        let some_block = self.context.append_basic_block(func, "filter_some");
        let none_block = self.context.append_basic_block(func, "filter_none");
        let predicate_true_block = self.context.append_basic_block(func, "filter_true");
        let predicate_false_block = self.context.append_basic_block(func, "filter_false");
        let continue_block = self.context.append_basic_block(func, "filter_continue");

        self.builder
            .build_conditional_branch(is_some, some_block, none_block)?;

        // Some block: extract value, call predicate
        self.builder.position_at_end(some_block);
        let some_value = self.extract_option_some_value(option_value, option_type)?;

        let some_type = self.nrc_type_to_llvm_type(&Some(*option_type.some_type.clone()))?;
        let some_temp_alloca = self.builder.build_alloca(some_type, "some_temp")?;
        self.builder.build_store(some_temp_alloca, some_value)?;

        let some_temp_var = format!("__filter_some_{}", location.line);
        let old_some_value = self
            .variable_map
            .insert(some_temp_var.clone(), some_temp_alloca.into());
        let old_some_type = self
            .variable_type_map
            .insert(some_temp_var.clone(), *option_type.some_type.clone());

        // Call predicate function
        let func_expr = arguments[0].clone();
        let some_var_expr = Expression::Variable(some_temp_var.clone());
        let call_expr = Expression::Call(CallExpr {
            callee: Box::new(func_expr),
            arguments: vec![some_var_expr],
            type_args: Vec::new(),
            resolved: None,
            location: *location,
        });

        let predicate_result = self.generate_expression(&call_expr)?;

        // Restore variable map
        if let Some(old) = old_some_value {
            self.variable_map.insert(some_temp_var.clone(), old);
        } else {
            self.variable_map.remove(&some_temp_var);
        }
        if let Some(old) = old_some_type {
            self.variable_type_map.insert(some_temp_var, old);
        } else {
            self.variable_type_map.remove(&some_temp_var);
        }

        // Check predicate result (should be bool/int)
        let predicate_bool = match predicate_result {
            BasicValueEnum::IntValue(int_val) => {
                let int_type = int_val.get_type();
                let bit_width = int_type.get_bit_width();
                let zero = if bit_width == 1 {
                    // Bool type (i1)
                    self.context.bool_type().const_int(0, false)
                } else {
                    // Regular int type
                    self.context.i32_type().const_int(0, false)
                };
                // Cast int_val to i32 if needed for comparison
                let int_val_for_compare = if bit_width == 1 {
                    // Bool type (i1) - extend to i32 for comparison
                    self.builder.build_int_z_extend(
                        int_val,
                        self.context.i32_type(),
                        "bool_extend",
                    )?
                } else if bit_width != 32 {
                    // Extend or truncate to i32
                    if bit_width < 32 {
                        self.builder.build_int_s_extend(
                            int_val,
                            self.context.i32_type(),
                            "int_extend",
                        )?
                    } else {
                        self.builder.build_int_truncate(
                            int_val,
                            self.context.i32_type(),
                            "int_truncate",
                        )?
                    }
                } else {
                    int_val
                };
                let zero_for_compare = if bit_width == 1 {
                    self.context.i32_type().const_int(0, false)
                } else {
                    zero
                };
                self.builder.build_int_compare(
                    inkwell::IntPredicate::NE,
                    int_val_for_compare,
                    zero_for_compare,
                    "predicate_check",
                )?
            }
            BasicValueEnum::FloatValue(_) => {
                return Err(CompilerError::internal(
                    "filter() predicate must return bool or int",
                ));
            }
            _ => {
                return Err(CompilerError::internal(
                    "filter() predicate must return bool or int",
                ))
            }
        };

        self.builder.build_conditional_branch(
            predicate_bool,
            predicate_true_block,
            predicate_false_block,
        )?;

        // Predicate true: keep Some
        self.builder.position_at_end(predicate_true_block);
        self.builder.build_unconditional_branch(continue_block)?;

        // Predicate false: convert to None
        self.builder.position_at_end(predicate_false_block);
        let new_none = self.create_option_value(1, None, option_type)?;
        self.builder.build_unconditional_branch(continue_block)?;

        // None block: keep None
        self.builder.position_at_end(none_block);
        let none_result = self.create_option_value(1, None, option_type)?;
        self.builder.build_unconditional_branch(continue_block)?;

        // Continue block: return Some (if predicate true), None (if predicate false or original None)
        self.builder.position_at_end(continue_block);
        let option_llvm_type = self.get_option_llvm_type(option_type)?;
        let phi = self.builder.build_phi(option_llvm_type, "filter_result")?;
        phi.add_incoming(&[
            (&option_value, predicate_true_block),
            (&new_none, predicate_false_block),
            (&none_result, none_block),
        ]);
        Ok(phi.as_basic_value())
    }

}
